// Optimized bidirectional copy utilities

use tokio::io::{AsyncRead, AsyncWrite, AsyncReadExt, AsyncWriteExt};
use std::io::Result;

const BUFFER_SIZE: usize = 32 * 1024; // 32KB buffer

pub async fn bidirectional_copy<A, B>(a: &mut A, b: &mut B) -> Result<(u64, u64)>
where
    A: AsyncRead + AsyncWrite + Unpin,
    B: AsyncRead + AsyncWrite + Unpin,
{
    let (mut a_read, mut a_write) = tokio::io::split(a);
    let (mut b_read, mut b_write) = tokio::io::split(b);

    let a_to_b = async {
        let mut buf = vec![0u8; BUFFER_SIZE];
        let mut total = 0u64;
        loop {
            let n = a_read.read(&mut buf).await?;
            if n == 0 {
                break;
            }
            b_write.write_all(&buf[..n]).await?;
            total += n as u64;
        }
        Ok::<u64, std::io::Error>(total)
    };

    let b_to_a = async {
        let mut buf = vec![0u8; BUFFER_SIZE];
        let mut total = 0u64;
        loop {
            let n = b_read.read(&mut buf).await?;
            if n == 0 {
                break;
            }
            a_write.write_all(&buf[..n]).await?;
            total += n as u64;
        }
        Ok::<u64, std::io::Error>(total)
    };

    tokio::try_join!(a_to_b, b_to_a)
}

pub async fn copy_with_callback<R, W, F>(
    reader: &mut R,
    writer: &mut W,
    mut callback: F,
) -> Result<u64>
where
    R: AsyncRead + Unpin,
    W: AsyncWrite + Unpin,
    F: FnMut(usize),
{
    let mut buf = vec![0u8; BUFFER_SIZE];
    let mut total = 0u64;

    loop {
        let n = reader.read(&mut buf).await?;
        if n == 0 {
            break;
        }
        writer.write_all(&buf[..n]).await?;
        callback(n);
        total += n as u64;
    }

    Ok(total)
}
